'''
Created on Mar 22, 2012

@author: pedro
'''

import Preprocessor as pre
import HDRGenerators
import HDRNormalizers
import Metaprocessing
import sys
import os
import time
import yappi #@UnresolvedImport
import numpy as np


from Network import OscillatorPyramidalLayer,OutputLayer
from Neurons import PyramidalNeuron,OutputNeuron,integrate_timecourse
from Utils import angle

from random import shuffle,randint
from multiprocessing import Process,cpu_count,current_process



def hdr_battery_heatmap(battery,filename):
    battery_arr = range(len(battery))    

    for key, value in battery.iteritems():
        battery_arr[key] = value
    
    pts = [range(len(battery_arr)) for x in range(len(battery_arr))]
 
    count = len(battery_arr) * (len(battery_arr)-1)
    total = 0.0
 
    for x in range(len(battery_arr)):
        for y in range(len(battery_arr)):
            diff = angle(battery_arr[x],battery_arr[y]) if x != y else 0.0
            total += diff
            pts[x][y] = diff
    
    for x in range(len(battery_arr)):
        for y in range(len(battery_arr)):
            if pts[x][y] == 0:
                pts[x][y] = (total)/(count)
    
    if filename != None:
        Metaprocessing.save_to_heatmap(pts,filename)
        
    return pts


def init_hdr_battery(hdr_file, config_params):
    battery = None
    try:
        if os.stat(hdr_file)[6]==0:
            raise Exception
        else:
            print "Loading battery from file...."
            battery = pre.load_battery_from_file(hdr_file)
            print "... loaded battery, found {0} HDRs".format(len(battery))
    except:
        print "Could not find battery file, creating a new one"
#        gen = HDRGenerators.UniformSigmoidHDRGenerator(config_params["usigmoid_intercept"],
#                                                       config_params["usigmoid_slope"],
#                                                       config_params["usigmoid_gaussian_mu"],
#                                                       config_params["usigmoid_gaussian_sigma"])
        gen = HDRGenerators.UniformSigmoidStarHDRGenerator(
                                               config_params["usigmoid_intercept"],
                                               config_params["usigmoid_slope"],
                                               config_params["usigmoid_gaussian_mu"],
                                               config_params["usigmoid_gaussian_sigma"])
        
#        gen = HDRGenerators.RandomizedNVectorHDRGenerator(1.00)
        
        norm = HDRNormalizers.FloorHDRNormalizer()
        battery = pre.generate_hdr_battery(config_params["num_hdrs"]+1, config_params["numPyrs"], gen, norm)
#        pre.save_battery_to_file(battery, hdr_file)
        
    return battery


def init_neural_net(num_outputs, config_params):
    pyramidals = []
    outputs = []
    for i in range(num_outputs):
        o1 = OutputNeuron(config_params["v0"],
                          config_params["Rm"],
                          config_params["Cm"],
                          config_params["tau_ref"],
                          config_params["Vth"],
                          config_params["V_spike"])
        outputs.append(o1)
        o1 = None
    
    for i in range(config_params["numPyrs"]):
        pyr1 = PyramidalNeuron(config_params["v0"],
                               config_params["Rm"],
                               config_params["Cm"],
                               config_params["tau_ref"],
                               config_params["Vth"],
                               config_params["V_spike"])
        pyramidals.append(pyr1)
        pyr1 = None

    pyrlayer = OscillatorPyramidalLayer(pyramidals, config_params["oscillator_freq"])
    outlayer = OutputLayer(pyramidals, 
                           outputs, 
                           config_params["stdp_decay_factor"], 
                           config_params["max_weight"], 
                           config_params["connectivity"])
    return (pyrlayer, outlayer)
    

def run_simulation(config_params, hdr_file, data_file):
    
    '''---Turn on Profiling if config--'''
    if config_params["profiling"]:
        yappi.start()
    
    '''---Initialize HDR Battery---'''  
    battery = init_hdr_battery(hdr_file, config_params)
    hdr_battery_angle_matrix = np.array(hdr_battery_heatmap(battery, None))
    np.savetxt("hdr_bat_{0}.txt".format(current_process().ident), hdr_battery_angle_matrix)
#    hdr_battery_heatmap(battery,"hdr_{0}_{1}.png".format(current_process().ident,time.strftime("%d_%b_%Y_%H-%M-%S")))
    
    
    '''---Initialize Neural Network---'''
    numOuts = config_params["numPyrs"]
    network = init_neural_net(numOuts, config_params)    
    pyrlayer = network[0]
    outlayer = network[1]
    
    
    hdr_baseline_outputs = {}
    hdr_posttraining_outputs = {}
    
    '''---Choose P+ HDRs Randomly---'''
    pplus_hdrs = []
    while len(pplus_hdrs) < config_params["num_pplus_hdrs"]:
        attempt = randint(0,len(battery)-1)
        if attempt in pplus_hdrs:
            continue
        else:
            pplus_hdrs.append(attempt)

    
    '''---Part 1: Establish Baseline Responses---'''
    
    
    """Each HDR is presented $num_regular_trials times at random"""
    
    trials_sequence = []
    for i in range(len(battery)):
        for j in range(config_params["num_regular_trials"]):
            trials_sequence.append(i)
    shuffle(trials_sequence)
    
    for hdr_num in trials_sequence:
        #Reset neuron's Vm to v0
        pyrlayer.reset_neurons()
        outlayer.reset_neurons()
        
        #Fire Pyramidal Layer first
        pi_layer_outputs = pyrlayer.fire(battery[hdr_num], config_params["dt"], config_params["time"])
        
        #Fire Output Layer Last
        output_layer_outputs = outlayer.fire(pi_layer_outputs, 
                                             config_params["dt"], 
                                             config_params["time"], 
                                             True, 
                                             config_params["Aplus"], 
                                             config_params["Aminus"], 
                                             config_params["stdp_tauplus"], 
                                             config_params["stdp_tauminus"])
        hdr_baseline_outputs[hdr_num] = sum([integrate_timecourse(output[1], config_params["v0"], config_params["time_stop"]) for output in output_layer_outputs])
    
    pyrlayer.reset_neurons()
    outlayer.reset_neurons()
    
    """Acquire Baseline Responses by Firing w/ Learning OFF"""
    for hdr_num in range(len(battery)):
        #Fire Pyramidal Layer first
        pi_layer_outputs = pyrlayer.fire(battery[hdr_num], config_params["dt"], config_params["time"])
        
        #Fire Output Layer Last
        output_layer_outputs = outlayer.fire(pi_layer_outputs, 
                                             config_params["dt"], 
                                             config_params["time"], 
                                             False, 
                                             config_params["Aplus"], 
                                             config_params["Aminus"], 
                                             config_params["stdp_tauplus"], 
                                             config_params["stdp_tauminus"])
        
        hdr_baseline_outputs[hdr_num] = sum([integrate_timecourse(output[1], config_params["v0"], config_params["time_stop"]) for output in output_layer_outputs])
        
        pyrlayer.reset_neurons()
        outlayer.reset_neurons()


    '''---Part 2: Perform Additional Training for P+ HDRs---'''
        
    for x in range(config_params["num_extra_trials"]):
        for hdr_num in pplus_hdrs:
            #Reset neuron's Vm to v0
            pyrlayer.reset_neurons()
            outlayer.reset_neurons()
            
            #Fire Pyramidal Layer first
            pi_layer_outputs = pyrlayer.fire(battery[hdr_num], config_params["dt"], config_params["time"])
            
            #Fire Output Layer Last
            output_layer_outputs = outlayer.fire(pi_layer_outputs, 
                                                 config_params["dt"], 
                                                 config_params["time"], 
                                                 True, 
                                                 config_params["Aplus"], 
                                                 config_params["Aminus"], 
                                                 config_params["stdp_tauplus"], 
                                                 config_params["stdp_tauminus"])
            hdr_posttraining_outputs[hdr_num] = sum([integrate_timecourse(output[1], config_params["v0"], config_params["time_stop"]) for output in output_layer_outputs])
        
    pyrlayer.reset_neurons()
    outlayer.reset_neurons()
    
    
    
    '''---Part 3: Acquire Post-Training Responses by Firing w/ Learning OFF---'''
    
    for hdr_num in range(len(battery)):
        #Fire Pyramidal Layer first
        pi_layer_outputs = pyrlayer.fire(battery[hdr_num], config_params["dt"], config_params["time"])
        
        #Fire Output Layer Last
        output_layer_outputs = outlayer.fire(pi_layer_outputs, 
                                             config_params["dt"], 
                                             config_params["time"], 
                                             False, 
                                             config_params["Aplus"], 
                                             config_params["Aminus"], 
                                             config_params["stdp_tauplus"], 
                                             config_params["stdp_tauminus"])
        
        hdr_posttraining_outputs[hdr_num] = sum([integrate_timecourse(output[1], config_params["v0"], config_params["time_stop"]) for output in output_layer_outputs])
        
        pyrlayer.reset_neurons()
        outlayer.reset_neurons()
        
    bminus = []
    bplus = []
    pminus = []
    pplus = []
    output_data = []
    
    for hdr_num in range(len(battery)):
        if hdr_num in pplus_hdrs:
            output_data.append((1,hdr_num,hdr_baseline_outputs[hdr_num],hdr_posttraining_outputs[hdr_num],hdr_posttraining_outputs[hdr_num]-hdr_baseline_outputs[hdr_num]))
            bplus.append((hdr_num,hdr_baseline_outputs[hdr_num]))
            pplus.append((hdr_num,hdr_posttraining_outputs[hdr_num]))
        else:
            output_data.append((0,hdr_num,hdr_baseline_outputs[hdr_num],hdr_posttraining_outputs[hdr_num],hdr_posttraining_outputs[hdr_num]-hdr_baseline_outputs[hdr_num]))
            bminus.append((hdr_num,hdr_baseline_outputs[hdr_num]))
            pminus.append((hdr_num,hdr_posttraining_outputs[hdr_num]))
        
    print "\nBaseline v.s. Post-Training for B-/P- HDRs:"    
    for index in range(len(bminus)):
        print "\tHDR #{0} B-: {1}".format(bminus[index][0],bminus[index][1])
        print "\tHDR #{0} P-: {1}".format(pminus[index][0],pminus[index][1])
        print "\tHDR #{0} Delta: {1}".format(pminus[index][0],pminus[index][1]-bminus[index][1])
        
    print "\nBaseline v.s. Post-Training for B+/P+ HDRs:"    
    for index in range(len(bminus)):
        print "\tHDR #{0} B+: {1}".format(bplus[index][0],bplus[index][1])
        print "\tHDR #{0} P+: {1}".format(pplus[index][0],pplus[index][1])
        print "\tHDR #{0} Delta: {1}".format(pplus[index][0],pplus[index][1]-bplus[index][1])
    
    print "Done"
    
    Metaprocessing.write_to_output_file(output_data, data_file)
    
    if config_params["profiling"]:
        yappi.print_stats()

"""--------------------------------------------------------------------------------------------------"""

if __name__ == '__main__':
    
    config_file = sys.argv[1]
    hdr_file = sys.argv[2]
    data_file = sys.argv[3]
    
    '''---Load Config---'''
    config_params = Metaprocessing.get_config_params(config_file)
        
    if config_params["multithreading"]:
        for i in range(config_params["num_iterations"]/cpu_count()):
            processes = []
            for i in range(cpu_count()):
                p = Process(target=run_simulation, args=(config_params, hdr_file, data_file))
                processes.append(p)
            
         
            for i in range(len(processes)):
                print "Starting Process {0} at time {1}...".format(i, time.asctime())
                processes[i].start()
                
            for i in range(len(processes)):
                processes[i].join()
                print "Process {0} re-joined at time {1}".format(i, time.asctime())  
    else:
        for i in range(config_params["num_iterations"]): 
            run_simulation(config_params, hdr_file, data_file)
            
    Metaprocessing.get_average_hdr_heatmap()
    Metaprocessing.calculate_results_from_output(data_file)
    Metaprocessing.log_cleanup()
    
    
    

    